import numpy as np
import scipy.sparse as sp
from scipy.sparse.linalg import spsolve
import os

def MFEABernoulli3Ddata(filename):
    Node = []
    Element = []
    Material = []
    Load = []
    Constr = []
    DF = []

    isNodeSection = False
    isElementSection = False
    isMaterialSection = False
    isLoadSection = False
    isDFSection = False
    isConstrSection = False

    with open(os.path.join("input", f"{filename}.txt"), 'r') as fid:
        for line in fid:
            line = line.strip()
            if not line:
                continue

            if line.startswith('*Node'):
                isNodeSection = True
                isElementSection = False
                isLoadSection = False
                isConstrSection = False
                isDFSection = False
                isMaterialSection = False
                continue
            elif line.startswith('*Element'):
                isNodeSection = False
                isElementSection = True
                isLoadSection = False
                isConstrSection = False
                isDFSection = False
                isMaterialSection = False
                continue
            elif line.startswith('*Material'):
                isNodeSection = False
                isElementSection = False
                isLoadSection = False
                isConstrSection = False
                isDFSection = False
                isMaterialSection = True
                continue
            elif line.startswith('*Load'):
                isNodeSection = False
                isElementSection = False
                isLoadSection = True
                isConstrSection = False
                isDFSection = False
                isMaterialSection = False
                continue
            elif line.startswith('*DF'):
                isNodeSection = False
                isElementSection = False
                isLoadSection = False
                isDFSection = True
                isConstrSection = False
                isMaterialSection = False
                continue
            elif line.startswith('*Constr'):
                isNodeSection = False
                isElementSection = False
                isLoadSection = False
                isConstrSection = True
                isDFSection = False
                isMaterialSection = False
                continue

            if isNodeSection:
                data = list(map(float, line.split(',')))
                if len(data) >= 4:
                    Node.append(data[1:])
            elif isElementSection:
                data = list(map(int, line.split(',')))
                if len(data) >= 3:
                    Element.append(data[1:])
            elif isMaterialSection:
                data = list(map(float, line.split(',')))
                if len(data) >= 6:
                    Material.append(data[:6])
            elif isLoadSection:
                data = list(map(float, line.split(',')))
                if len(data) >= 3:
                    Load.append([int(data[0]), int(data[1]), data[2]])
            elif isDFSection:
                data = list(map(float, line.split(',')))
                if len(data) >= 4:
                    DF.append(data)
            elif isConstrSection:
                data = list(map(float, line.split(',')))
                if len(data) >= 3:
                    Constr.append([int(data[0]), int(data[1]), data[2]])
    return np.array(Node), np.array(Element), np.array(Material), np.array(Load), np.array(Constr)


def Bernoulli3DFEA(Node, Element, Material, force, constrain):
    Dof = 6
    EDof = Dof * 2
    Dofs = Dof * Node.shape[0]
    EleCount = Element.shape[0]
    KKG = sp.lil_matrix((Dofs, Dofs))
    FFG = sp.lil_matrix((Dofs, 1))
    U = sp.lil_matrix((Dofs, 1))
    x = Node[:, 0]
    y = Node[:, 1]
    z = Node[:, 2]
    BarLength = BarsLength(x, y, z, Element)
    print(BarLength)
    for iEle in range(EleCount):
        n1, n2 = Element[iEle] - 1
        R = CoordTransform([x[n1], x[n2]], [y[n1], y[n2]], [z[n1], z[n2]], BarLength[iEle])
        ElementStiffnessMatrix = Bernoulli3DElementKe(Material[0], R, BarLength[iEle], 1)
        n1_dofs = np.arange(n1 * Dof, (n1 + 1) * Dof)
        n2_dofs = np.arange(n2 * Dof, (n2 + 1) * Dof)
        ElementNodeDOF = np.concatenate((n1_dofs, n2_dofs))
        for i in range(len(ElementNodeDOF)):
            for j in range(len(ElementNodeDOF)):
                KKG[ElementNodeDOF[i], ElementNodeDOF[j]] += ElementStiffnessMatrix[i, j]



    if force.size:
        for i in range(len(force)):
            m = int(force[i, 0])
            n = int(force[i, 1])
            value = force[i, 2]
            FFG[Dof * (m - 1) + n - 1, 0] = value

    for i in range(len(constrain)):
        n = int(constrain[i, 0])  # 节点号
        d = int(constrain[i, 1])  # 局部自由度号
        m = (n - 1) * Dof + d - 1  # 全局自由度索引位置
        FFG = FFG - constrain[i, 2] * KKG[:, m]  # 载荷列阵均减去相应元素
        FFG[m, 0] = constrain[i, 2]  # 载荷列阵相应位置更改为已知位移
        KKG[:, m] = 0  # 刚度矩阵列置零
        KKG[m, :] = 0  # 刚度矩阵行置零
        KKG[m, m] = 1.0  # 刚度矩对角元素阵置 1

    print(KKG)

    # 转换为全矩阵
    dense_matrix = KKG.toarray()
    KKGG = np.linalg.inv(dense_matrix)


    U = spsolve(KKG.tocsc(), FFG)
    # 计算应变和应力
    Strain, Stress = ComputeStrainStress(Node, Element, U, Material)
    return U, Strain, Stress

def ComputeStrainStress(Node, Element, U, Material):
    Dof = 6
    Strain = []
    Stress = []
    for ele in Element:
        n1, n2 = ele - 1
        Ue = np.concatenate((U[n1*Dof:(n1+1)*Dof], U[n2*Dof:(n2+1)*Dof]))
        Le = np.linalg.norm(Node[n2] - Node[n1])
        e = (Ue[6] - Ue[0]) / Le
        E = Material[0][0]
        stress = E * e
        Strain.append(e)
        Stress.append(stress)
    return np.array(Strain), np.array(Stress)

def BarsLength(x, y, z, ele):
    return np.sqrt((x[ele[:, 1] - 1] - x[ele[:, 0] - 1]) ** 2 +
                   (y[ele[:, 1] - 1] - y[ele[:, 0] - 1]) ** 2 +
                   (z[ele[:, 1] - 1] - z[ele[:, 0] - 1]) ** 2)


def CoordTransform(x, y, z, L):
    x1, x2 = x
    y1, y2 = y
    z1, z2 = z
    if x1 == x2 and y1 == y2:
        if z2 > z1:
            Lambda = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
        else:
            Lambda = np.array([[0, 0, -1], [0, 1, 0], [1, 0, 0]])
    else:
        CXx = (x2 - x1) / L
        CYx = (y2 - y1) / L
        CZx = (z2 - z1) / L
        D = np.sqrt(CXx * CXx + CYx * CYx)
        CXy = -CYx / D
        CYy = CXx / D
        CZy = 0
        CXz = -CXx * CZx / D
        CYz = -CYx * CZx / D
        CZz = D
        Lambda = np.array([[CXx, CYx, CZx],
                           [CXy, CYy, CZy],
                           [CXz, CYz, CZz]])
    R = np.block([
        [Lambda, np.zeros((3, 9))],
        [np.zeros((3, 3)), Lambda, np.zeros((3, 6))],
        [np.zeros((3, 6)), Lambda, np.zeros((3, 3))],
        [np.zeros((3, 9)), Lambda]
    ])
    return R


def Bernoulli3DElementKe(prop, R, L, icoord):
    E, A, Iy, Iz, G, J = prop
    k1 = E * A / L
    k2 = 12 * E * Iz / (L ** 3)
    k3 = 6 * E * Iz / (L ** 2)
    k4 = 4 * E * Iz / L
    k5 = 2 * E * Iz / L
    k6 = 12 * E * Iy / (L ** 3)
    k7 = 6 * E * Iy / (L ** 2)
    k8 = 4 * E * Iy / L
    k9 = 2 * E * Iy / L
    k10 = G * J / L
    ke = np.array([
        [k1, 0, 0, 0, 0, 0, -k1, 0, 0, 0, 0, 0],
        [0, k2, 0, 0, 0, k3, 0, -k2, 0, 0, 0, k3],
        [0, 0, k6, 0, -k7, 0, 0, 0, -k6, 0, -k7, 0],
        [0, 0, 0, k10, 0, 0, 0, 0, 0, -k10, 0, 0],
        [0, 0, -k7, 0, k8, 0, 0, 0, k7, 0, k9, 0],
        [0, k3, 0, 0, 0, k4, 0, -k3, 0, 0, 0, k5],
        [-k1, 0, 0, 0, 0, 0, k1, 0, 0, 0, 0, 0],
        [0, -k2, 0, 0, 0, -k3, 0, k2, 0, 0, 0, -k3],
        [0, 0, -k6, 0, k7, 0, 0, 0, k6, 0, k7, 0],
        [0, 0, 0, -k10, 0, 0, 0, 0, 0, k10, 0, 0],
        [0, 0, -k7, 0, k9, 0, 0, 0, k7, 0, k8, 0],
        [0, k3, 0, 0, 0, k5, 0, -k3, 0, 0, 0, k4]
    ])
    print(ke)
    if icoord == 1:
        Ke = R.T @ ke @ R
    elif icoord == 2:
        Ke = ke
    return Ke



def Bernoulli3D2vtk(Node, Element, U, Strain, Stress, Scale, filename):
    UXYZ = np.zeros((Node.shape[0], 3))
    UXYZ[:, 0] = U[0::6]
    UXYZ[:, 1] = U[1::6]
    UXYZ[:, 2] = U[2::6]
    Node = Node + Scale * UXYZ

    num_nodes = Node.shape[0]
    num_elements = Element.shape[0]

    filePath1 = os.path.join('output', f'{filename}-Scale-{Scale}.vtk')
    with open(filePath1, 'w') as vtk_file:
        # 写入 VTK 文件头
        vtk_file.write('# vtk DataFile Version 3.0\n')
        vtk_file.write('Finite Element Results\n')
        vtk_file.write('ASCII\n')
        vtk_file.write('DATASET UNSTRUCTURED_GRID\n')

        # 写入节点坐标
        vtk_file.write(f'POINTS {num_nodes} float\n')
        for i in range(num_nodes):
            vtk_file.write(f'{Node[i, 0]} {Node[i, 1]} {Node[i, 2]}\n')

        # 写入单元信息
        vtk_file.write(f'CELLS {num_elements} {num_elements * 3}\n')
        for i in range(num_elements):
            vtk_file.write(f'2 {Element[i, 0] - 1} {Element[i, 1] - 1}\n')

        # 写入单元类型
        vtk_file.write(f'CELL_TYPES {num_elements}\n')
        for i in range(num_elements):
            vtk_file.write('3\n')  # 2 节点线元素

        # 写入位移数据
        vtk_file.write(f'POINT_DATA {num_nodes}\n')
        vtk_file.write('SCALARS U float 3\n')
        vtk_file.write('LOOKUP_TABLE default\n')
        for i in range(num_nodes):
            vtk_file.write(f'{U[i * 6]} {U[i * 6 + 1]} {U[i * 6 + 2]}\n')

        # 写入应变数据
        vtk_file.write(f'CELL_DATA {num_elements}\n')
        vtk_file.write('SCALARS Strain float 1\n')
        vtk_file.write('LOOKUP_TABLE default\n')
        for i in range(num_elements):
            vtk_file.write(f'{Strain[i]}\n')

        # 写入应力数据
        vtk_file.write('SCALARS Stress float 1\n')
        vtk_file.write('LOOKUP_TABLE default\n')
        for i in range(num_elements):
            vtk_file.write(f'{Stress[i]}\n')



if __name__ == '__main__':
    filename = 'liuchang'
    Scale = 1
    Node, Element, Material, Load, Constr = MFEABernoulli3Ddata(filename)
    U, Strain, Stress = Bernoulli3DFEA(Node, Element, Material, Load, Constr)
    Bernoulli3D2vtk(Node, Element, U, Strain, Stress, Scale, filename)